#include "dfruntime.h"
#include "dfruntime_private.h"
#include "global_queue.h"
#include "memory_management.h"
#include "asyncpages.h"
#include <linux/unistd.h>
#include <time.h>


static struct glob_task_bag *intermediate_buffer = NULL;
static struct glob_task_bag *global_queue = NULL;
static struct global_register * global_register;

#define MAX_NODE 256

// External queue element.
struct glob_task_bag {
    struct glob_task_bag *next;
    int count;
    struct frame_struct * data[GLOBQUEUE_ELT_TASK_AMOUNT];
};

struct global_register {
    int noden;
    void * node_pointers[MAX_NODE];
    bool no_work[MAX_NODE]; // For termination.
};

static int __inited  = 0;


void gqueue_print_queue() {
    sync_safe_acquire(global_queue, sizeof(*global_queue), true);
    DEBUG("Printing queue:");
    struct glob_task_bag *cur = global_queue;
    int z = 0;
    while ( cur != NULL ) {
        DEBUG("Bag :: %d : %d tasks",z++, cur->count);
        for ( int i = 0; i < cur->count; ++i ) {
            DEBUG("bag has %p", cur->data[i]);
        }
        cur = cur->next;
    }

    safe_release_page(global_queue);

}

// TODO : fix races b/w adding to the global queue and stealing ( for cleanup ).
//
void gqueue_init_pre_fork() {

  // Allocate memory for global register.
        void * reg = dfrt_mem_malloc( sizeof(struct global_register) );
        global_register = reg;
        dfrt_mem_set_global(global_register);
        //dfrt_mem_blocking_exclusive_acquire(global_register, sizeof(struct global_register) );
        memset( global_register,0,sizeof(struct global_register));
        global_register->noden = dfrt_get_num_nodes();
        //safe_release_page( global_register );



}

// Must be done once per node.
void gqueue_init_last() {

        if ( ! __inited ) {
            // Allocates GlobalQueue, as well as the Intermediate buffer.
            intermediate_buffer = calloc( 1, sizeof( struct glob_task_bag) );
            global_queue = dfrt_mem_malloc( sizeof(struct glob_task_bag));
            sync_safe_acquire( global_queue, sizeof( struct glob_task_bag), true);
            global_queue->next = NULL;
            safe_release_page( global_queue );

            // Registers the queue into the global register.
            global_register = dfrt_mem_get_global();
            sync_safe_acquire(global_register, sizeof(struct global_register), true );
            global_register->node_pointers[dfrt_get_node_num()] = global_queue;

            safe_release_page( global_register );



            __inited = true;
        } else {
            FATAL("Double initialization of global queue.");
        }

}

void gqueue_set_termination_bit( bool val ) {
    sync_safe_acquire( global_register, sizeof( struct global_register ), true);
    global_register->no_work[dfrt_get_node_num()] = val;
    safe_release_page( global_register);
}


bool gqueue_check_termination( )  {
    sync_safe_acquire( global_register, sizeof( struct global_register), false);
    for ( int i = 0; i < global_register->noden; ++ i ) {
        if ( global_register->no_work[i] == false ) {
            DEBUG("Node %d not terminated !",i);
           safe_release_page( global_register);
           return false;
        }
    }
    safe_release_page( global_register);

    return true;
}

void gqueue_append( struct frame_struct *fp) {
    #pragma omp critical (gqueue)
    {
	// If interbuffer is full : add it to global queue.
	if ( intermediate_buffer->count == GLOBQUEUE_ELT_TASK_AMOUNT ) {
            gqueue_export_interbuf();
	}

	// Append the task effectively.
	intermediate_buffer->data[ intermediate_buffer->count++ ] = fp;
    }

}

void gqueue_export_interbuf() {
    _DEBUG("Exporting interbuf");


        sync_safe_acquire( global_queue, sizeof( struct glob_task_bag), true);
	// Remove empty bags :
	struct glob_task_bag *pre = global_queue;
	struct glob_task_bag *next = global_queue->next;
	while ( next != NULL) {
            if ( next->count == 0 ) {
                // We must free it.
                pre->next = next->next;
                //dfrt_mem_free(next);
                next = pre->next;
            } else {
                pre = next;
                next = pre->next;
            }
	}


	if ( global_queue->count == 0 ) {
            // Then we copy the pointer and replace the first bag.
            next = global_queue->next;
            memcpy( global_queue, intermediate_buffer, sizeof(struct glob_task_bag));
            global_queue->next = next;

	} else {
            // Then we add it to the queue...
            struct glob_task_bag * newbag = dfrt_mem_malloc( sizeof( struct glob_task_bag ));
            sync_safe_acquire( newbag, sizeof( struct glob_task_bag), 1);
            memcpy( newbag, intermediate_buffer, sizeof(struct glob_task_bag));

            newbag->next = global_queue->next;
            global_queue->next = newbag;
            safe_release_page( newbag );

	}
        safe_release_page( global_queue );
        //gqueue_print_queue();
	// Then the intermediate buffer can be erased :
        memset( intermediate_buffer, 0, sizeof( struct glob_task_bag) );

}

// Returns the number of migrated tasks.
int gqueue_reacquire_interbuf() {
    _DEBUG("Reacquiring Interbuf");
    int retval = 0;

    #pragma omp critical (gqueue)
    {
	for ( int i = 0; i < intermediate_buffer->count; ++i ) {
            if ( intermediate_buffer->data[i] != NULL ) {
                schedule_forlocal( intermediate_buffer->data[i]);
                retval++;
            }
	}
	intermediate_buffer->count = 0;
    }
    return retval;
}


int gqueue_steal_local( int minimum ) {
    int stolen = 0;
    #pragma omp critical (gqueue)
    {
        // TODO : needs to keep global queue acquired to avoid races.
        struct glob_task_bag *current = global_queue;
        sync_safe_acquire( current, sizeof( struct glob_task_bag ), true );
        sync_safe_acquire( global_queue, sizeof( struct glob_task_bag ), true );

        while ( stolen < minimum ) {
            for ( int i = 0; i < current->count; ++i ) {
                if ( current->data[i] != NULL ) {
                    schedule_forlocal( current->data[i] );
                    stolen++;
                    current->data[i] = NULL;
                }
            }
            current->count = 0;
            if ( current->next == NULL ) {
                break;
            }
            struct glob_task_bag * next = current->next;
            safe_release_page( current );
            current = next;
            sync_safe_acquire(current, sizeof( struct glob_task_bag), true);
        }

        safe_release_page(current);
        safe_release_page(global_queue);
    }
    DEBUG("Stealing from local : %d tasks", stolen);

    return stolen;
}

// Steal some tasks : returns the number of tasks stolen.
int gqueue_steal_tasks( int minimum ) {
    int stolen = 0;
    int tries = 0;

#pragma omp critical (gqueue)
    {
        // Mark the node as potentially terminating.
        gqueue_set_termination_bit(true);
        while ( true ) {
            tries ++;
            // Find a target.
            sync_safe_acquire( global_register, sizeof( struct global_register ), true);
            if ( global_register->noden <= 1 ) {
                // Nothing to steal.
                DEBUG( "One node registered at most : nothing to steal");
            }

            // Pick a satisfactory node :
            int chosen = -1;
            int randtries = 1;
            do {
                chosen = rand() % global_register->noden;
                randtries ++;
                if ( randtries > 500 ) {
                    DEBUG( "Too many tries !!! Giving up.");
                    break;
                }
                // If target is not started, reset and yield the page:
                //randtries = 0;
                safe_release_page( global_register );
                struct timespec z;
                z.tv_sec = 0;
                if ( randtries % 10 == 0 ) {
                    z.tv_nsec = 10 * 1000 * 1000 ;
                } else {
                    z.tv_nsec = 1 * 1000 * 1000 ;

                }
                nanosleep(&z,NULL);
                sync_safe_acquire( global_register, sizeof(struct global_register), true);

            } while ( chosen == dfrt_get_node_num()
                      || global_register->node_pointers[chosen] == NULL
                      );
            // Steal work there :
            struct glob_task_bag * queue_first = global_register->node_pointers[chosen];
            struct glob_task_bag * queue = global_register->node_pointers[chosen];
            CFATAL( queue == NULL, "No queue registered.");
            safe_release_page( global_register );
            sync_safe_acquire( queue_first, sizeof( struct glob_task_bag ), true );
            sync_safe_acquire( queue, sizeof( struct glob_task_bag ), true );

            // Position to a non empty zone
            while ( queue->count == 0 && queue->next != NULL ) {
                struct glob_task_bag * tmp = queue-> next;
                safe_release_page( queue );
                queue = tmp;
                sync_safe_acquire( queue, sizeof( struct glob_task_bag ), true );

            }

            // Transfer the tasks.
            for ( int i = 0; i < queue->count; ++i ) {
                schedule_forlocal(queue->data[i]);
                queue->data[i] = NULL;
                stolen ++;
            }

            queue->count = 0;

            safe_release_page( queue );
            safe_release_page( queue_first );

            if ( stolen != 0 ) {
                gqueue_set_termination_bit(false);
                break;
            } else  {
                if ( gqueue_check_termination() ) {
                    // Program terminated successfully.
                    DEBUG("All do stealing : EXIT");
                    exit_group(EXIT_SUCCESS);

                }
            }



        }
    }
    return stolen;

}

